package com.hyzs.utils;

import org.elasticsearch.action.admin.indices.delete.DeleteIndexRequest;
import org.elasticsearch.action.bulk.BulkProcessor;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.support.master.AcknowledgedResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.indices.CreateIndexRequest;
import org.elasticsearch.client.indices.CreateIndexResponse;
import org.elasticsearch.client.indices.GetIndexRequest;
import org.elasticsearch.common.unit.ByteSizeUnit;
import org.elasticsearch.common.unit.ByteSizeValue;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.common.xcontent.XContentType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.nio.charset.Charset;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

public class EsBulkFunctions {

    private static final Logger logger = LoggerFactory.getLogger(EsBulkFunctions.class);


    // read index setting file(.json), create index mappings and settings.
    public void rebuildIndex(RestHighLevelClient client, String indexName, String indexConfigFile) throws Exception {
        GetIndexRequest getRequest = new GetIndexRequest(indexName);
        boolean exists = client.indices().exists(getRequest, RequestOptions.DEFAULT);
        if(exists){
            logger.info("index {} exists...", indexName);
            DeleteIndexRequest delRequest = new DeleteIndexRequest(indexName);
            AcknowledgedResponse deleteIndexResponse = client.indices().delete(delRequest, RequestOptions.DEFAULT);
            logger.info("delete {} index response: {}",indexName, deleteIndexResponse.isAcknowledged());
        }

        CreateIndexRequest createIndexRequest = new CreateIndexRequest(indexName);
        Path file = Paths.get(ClassLoader.getSystemResource(indexConfigFile).toURI());
        List<String> judgement_json = Files.newBufferedReader(file, Charset.forName("UTF-8"))
                .lines()
                .collect(Collectors.toList());
        String setting_string = String.join("", judgement_json);

        createIndexRequest.source(setting_string, XContentType.JSON);
        CreateIndexResponse createIndexResponse = client.indices().create(createIndexRequest, RequestOptions.DEFAULT);
        logger.info("create {} index: {}",indexName, createIndexResponse.isAcknowledged());
    }


    public void bulkInsert(RestHighLevelClient client, Connection mysqlConn,
                           String indexName, int startIndex, int endIndex) throws Exception {
        Statement statement = mysqlConn.createStatement();
        int batch = 5000;

        BulkProcessor bulkProcessor = BulkProcessor.builder(
                (request, bulkListener) -> client.bulkAsync(request, RequestOptions.DEFAULT, bulkListener),
                listener)
                .setBulkActions(1000)
                .setBulkSize(new ByteSizeValue(5, ByteSizeUnit.MB))
                .setFlushInterval(TimeValue.timeValueSeconds(5))
                .setConcurrentRequests(3)
                .build();

        for (int index = startIndex; index < endIndex; index++) {
            ResultSet resultSet = statement.executeQuery("SELECT * FROM judicial_document LIMIT " + index * batch + "," + batch);
            while (resultSet.next()) {
                String doc_id = StringUtils.trimBlank(resultSet.getString(1));
                String title = StringUtils.trimBlank(resultSet.getString(3));
                // trim relative law texts
                String tmp_content = StringUtils.trimBlank(resultSet.getString(4));
                String content = StringUtils.trimJudgementContent(tmp_content);
                IndexRequest indexRequest = new IndexRequest(indexName)
                        .id(doc_id)
                        .source(XContentType.JSON, "title", title, "content", content);
                bulkProcessor.add(indexRequest);
            }
            // synchronous execution
            //BulkResponse bulkResponse = client.bulk(bulkRequest, RequestOptions.DEFAULT);
        }

        bulkProcessor.flush();
        bulkProcessor.awaitClose(10L, TimeUnit.MINUTES);
        bulkProcessor.close();

    }

    private BulkProcessor.Listener listener = new BulkProcessor.Listener() {
        @Override
        public void beforeBulk(long executionId, BulkRequest request) {
            logger.info("Executing bulk [{}] with {} requests",
                    executionId, request.numberOfActions());
        }

        @Override
        public void afterBulk(long executionId, BulkRequest request,
                              BulkResponse response) {
            if (response.hasFailures()) {
                logger.warn("Bulk [{}] executed with failures", executionId);
            } else {
                logger.info("Bulk [{}] completed in {} milliseconds",
                        executionId, response.getTook().getMillis());
            }
        }

        @Override
        public void afterBulk(long executionId, BulkRequest request,
                              Throwable failure) {
            logger.error("Failed to execute bulk", failure);
        }
    };


    public static void main(String[] args) throws Exception{
        String indexName = "judgement_v3";
        String indexConfigFile = "JudgementSettings_v3.json";
        MySQLConnectionUtils connUtils = new MySQLConnectionUtils();
        connUtils.setUrl("jdbc:mysql://172.16.18.28:3306/judicial_documents?characterEncoding=utf-8");
        connUtils.setUser("guest");
        connUtils.setPassword("qwer1234");
        Connection mysqlConn = connUtils.getConnection();
        EsConnUtils esConnUtils = new EsConnUtils();
        esConnUtils.setUrl("172.16.119.31");
        esConnUtils.setPort(9200);
        RestHighLevelClient client = esConnUtils.getClient();

        EsBulkFunctions esFunctions = new EsBulkFunctions();
        //esFunctions.rebuildIndex(client, indexName, indexConfigFile);

        esFunctions.bulkInsert(client, mysqlConn, indexName, 100, 150);
        mysqlConn.close();
        esConnUtils.closeClient();
    }


}
